from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import StructType, StringType, IntegerType

if __name__ == '__main__':
    spark = SparkSession.builder. \
        appName("test"). \
        master("local[*]"). \
        config("spark.sql.shuffle.partitions", 2). \
        getOrCreate()

    """
    spark.sql.shuffle.partitions 参数指的是，在sql计算中，shuffle算子阶段默认的分区数是200个，
    对于集群模式来说，200个默认也算比较合适
    如果在Local下运行，200个很多，在调度上会带来额外的损耗
    所以在Local下建议修改比较低 比如2、4、10均可
    这个参数和Spark RDD中设置并行的参数 是相互独立的
    """

    # TODO 案例：电源评分
    # 数据下载地址： http://files.grouplens.org/datasets/movielens/
    # 下面 ml-100k/u.data

    schema = StructType().add("user_id", StringType()).add("movie_id", IntegerType()).add("rank", IntegerType()).add(
        "ts", StringType())
    df = spark.read.format("csv").option("sep", "\t").option("header", False).option("encoding", "utf-8").schema(
        schema).load("/Users/cdhuangchao3/tmp/spark_demo/u.data")

    # TODO 1 用户平均分计算
    df.groupBy("user_id") \
        .avg("rank") \
        .withColumnRenamed("avg(rank)", "avg_rank") \
        .withColumn("avg_rank", F.round("avg_rank", 2)) \
        .orderBy("avg_rank", ascending=False) \
        .show()
    # TODO 2 电影平均分查询
    df.createOrReplaceTempView("movie")
    spark.sql("""
    select movie_id, round(avg(rank), 2) as avg_rank from movie group by movie_id order by avg_rank desc
    """)

    # TODO 3 查询大于平均分的电影数量 # Row
    print("大于平均分电影的数量是", df.where(df['rank'] > df.select(F.avg(df['rank'])).first()['avg(rank)']).count())

    # TODO 4 查询高分电影中（>3）打分次数最多的用户，此人打分的平均分
    user_id = df.where("rank > 3"). \
        groupBy("user_id"). \
        count(). \
        withColumnRenamed("count", "cnt"). \
        orderBy("cnt", ascending=False). \
        limit(1). \
        first()['user_id']

    # 计算这个人的打分平均分
    df.filter(df['user_id'] == user_id). \
        select(F.round(F.avg("rank"), 2)).show()

    # TODO 5 查询每个用户的平均打分、最低打分、最高打分
    df.groupBy("user_id"). \
        agg(
        F.round(F.avg("rank"), 2).alias("avg_rank"),
        F.min("rank").alias("min_rank"),
        F.max("rank").alias("max_rank")
    ).show()

    # TODO 6 查询评分超过100次的电影，的平均分 排名TOP10
    # df2: DataFrame = df.groupBy("movie_id").count()
    # print(df2.withColumnRenamed("count", "cnt").orderBy("cnt", ascending=False).take(10))
    df.groupBy("movie_id").agg(
        F.count("movie_id").alias("cnt"),
        F.round(F.avg("rank"), 2).alias("avg_rank")
    ).where(F.col("cnt") > 100).orderBy("avg_rank", ascending=False).limit(10).show()

"""
1. agg: 它是GroupData对象的API，作用是 在里面可以写多个聚合
2. alias： 它是Column对象的API，可以针对一个列 进行改名
3. withColumnRenamed： 它是DataFrame的API，可以对DF中的列进行改名，一次改一个列，改多个列 可以链式调用
4. orderBy: DataFrame的API，进行排序，参数1是被排序的列，参数2是 升序（True）或降序 False
5. first： DataFrame的API，取出DF的第一行数据，返回值结果是Row对象
# Row对象就是一个数组，你可以通过row['列名'] 来取出当前行中，某一列的具体数值，返回值不再是DF 或者GroupData 或者Column
"""
